import os
from openai import OpenAI


def send_message(client, model, messages):
    """
    发送消息并获取流式响应。

    Args:
        client (OpenAI): OpenAI客户端。
        model (str): 使用的模型名称。
        messages (list): 对话消息列表。

    Returns:
        str: 助手的完整回复。
    """
    response = client.chat.completions.create(
        model=model,  # 模型名称
        messages=messages,  # 对话消息
        temperature=0.3,  # 温度
        stream=True,  # 开启流式响应
    )
    response_content = ""
    for chunk in response:
        chunk_message = chunk.choices[0].delta
        if chunk_message.content:
            response_content += chunk_message.content
            print(chunk_message.content, end='', flush=True)
    return response_content


def main():
    """
    主函数。
    """
    # 获取API密钥
    api_key = os.getenv("KIMI_API_KEY")
    if not api_key:
        raise ValueError("API密钥未找到。请确保已设置环境变量KIMI_API_KEY。")

    # 设置模型名称
    model = "moonshot-v1-8k"  # 模型名称

    # 初始化OpenAI客户端
    client = OpenAI(
        api_key=api_key,   # API密钥
        base_url="https://api.moonshot.cn/v1",   # 基础URL
    )

    # 初始对话历史
    messages = [
        {
            "role": "system",
            "content": "你是专业的MySQL助手，根据我的需求回答问题。",   # 系统提示
        },
        {"role": "user", "content": "你好"},
    ]

    while True:
        # 发送消息并获取响应
        response_content = send_message(client, model, messages)

        # 将响应添加到对话历史
        messages.append({"role": "assistant", "content": response_content})

        # 获取用户输入
        user_input = input("\n你: ")
        if user_input.lower() in ["exit", "quit"]:
            break

        # 将用户输入添加到对话历史
        messages.append({"role": "user", "content": user_input})


if __name__ == "__main__":
    main()
